# This module contains the casidi version of the finite difference model for the 1D Richards Equation (RE). 
# The 1D RE is expressed as C(h)dh/dt = d/dz(K(h)d/dz(h+z) where the terms have the following meanings
# C(h) -  differential water capacity/ capillary capacity
# h is -  the soil water pressure head
# K(h) -  is the unsaturated hydraulic conductivity
# z is -  the vertival coordinate directed upward

# Assumptions
# 1. The field is homogeneous (i.e same soil type) and no slope exists.
# 2. This module does not include the crop model.

# Boundary Conditions
# 1. Top BC    - Flux boundary condition [dh/dz= -1 - irrigation_amount/K(h)]
# 2. Bottom BC - Free Drainage [d(h+z)/dz= 1]

# The dependence of C, and K on h can be found in the van_GenucthenNumpy module.



import numpy as np 
from casadi import *
from Parameters1D_casadi import *
import van_Genuchten1D_casadi as vg 


def RichardsPolar_1D(head,irrigAmount):

    totalDepth, axialNodes, axialDistance, totalNodes = spatialVariables_1D()  #Spatial parameters
    
    
    coordinates  = generatingCoordinates_1D(axialNodes)
    
    dhdt = SX.zeros(totalNodes)

    for i in range(0,totalNodes):
        currentState = head[i]
        currentNode  = coordinates[i]
        curAxialNode = currentNode[-1] #The last coordinate represents the axial node

        if curAxialNode == 0: #This represents the bottom of the field

            stateSouth = currentState
            stateNorth = head[i + 1]
            hydCondNorth = vg.hydraulicConductivityApprox(currentState, stateNorth)
            hydCondSouth = vg.hydraulicConductivityApprox(currentState, stateSouth)
            fluxNorth = hydCondNorth*(((stateNorth   - currentState)/axialDistance) + 1)
            fluxSouth = hydCondSouth*(((currentState - stateSouth)/axialDistance) + 1)

            axialComponent = (1/axialDistance) * (fluxNorth - fluxSouth)

        elif curAxialNode == axialNodes - 1: # This represents the surface of the field where the irrigation is applied

            stateNorth = 0
            stateSouth = head[i - 1]
                        
            hydCondNorth = 0
            hydCondSouth =  vg.hydraulicConductivityApprox(currentState, stateSouth)
                    
            fluxNorth = -irrigAmount
            fluxSouth = hydCondSouth*(((currentState-stateSouth)/axialDistance) + 1)

            axialComponent = (1/axialDistance)*(fluxNorth - fluxSouth)

        
        else: #Interior nodes

            stateNorth = head[i + 1]
            stateSouth = head[i - 1]
            hydCondNorth = vg.hydraulicConductivityApprox(currentState, stateNorth)
            hydCondSouth = vg.hydraulicConductivityApprox(currentState, stateSouth)
            fluxNorth = hydCondNorth*(((stateNorth   - currentState)/axialDistance) + 1)
            fluxSouth = hydCondSouth*(((currentState - stateSouth)/axialDistance) + 1)

            axialComponent = (1/axialDistance)*(fluxNorth - fluxSouth)

        spatialApprox = (1/vg.capillaryCapacity(currentState))*(axialComponent)
        dhdt[i] = spatialApprox

    DhDt = dhdt
    return DhDt


def volMoistureAllNodes_1D(head):
    
    totalDepth, axialNodes, axialDistance, totalNodes = spatialVariables_1D()
    volMoistureInit = vg.volumetricMoisture(head)
    cMatrix = CMatrixAllMeasurements(totalNodes)
    volMoistureFin = mtimes(cMatrix, volMoistureInit)
    return volMoistureFin


##For this model,the cMatrix is defined as follows
def cMatrix_1D(axialNodes):

    cMatrix = np.zeros(axialNodes)

    for i in range(axialNodes):
        cMatrix[i] = (1/axialNodes)
        
    cMatrix = np.reshape(cMatrix, (1,16))

    return cMatrix



def volMoistureSelected_1D(cMatrix, head):
    volMoistureInit = vg.volumetricMoisture(head)
    volMoistureFin = mtimes(cMatrix, volMoistureInit)
    return volMoistureFin